from langchain import hub
from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain_community.agent_toolkits import create_sql_agent
from langchain_core.tools import Tool
from langchain_openai import ChatOpenAI


def get_response(key, input_text, db):
    model = ChatOpenAI(model="gpt-3.5-turbo",
                       openai_api_key=key,
                       openai_api_base="https://api.aigc369.com/v1",
                       temperature=0)

    db_agent_executor = create_sql_agent(
        db=db, llm=model, agent_type="openai-tools", verbose=True, handle_parsing_errors=True
    )

    def create_react_agent(db_agent_executor):
        tools = [
            Tool(
                name="数据库查询工具",
                description="""
                    当你本身无法获取到问题追定的数据信息的时候，你可以使用此数据库查询工具，将用户的问题转为SQL后，进行查询操作，最后工具会返回一个数据查询结果
                """,
                func=db_agent_executor.invoke
            ),
            Tool(
                name="数据库查询工具",
                description="""
                    当你本身无法获取到问题追定的数据信息的时候，你可以使用此数据库查询工具，将用户的问题转为SQL后，进行查询操作，最后工具会返回一个数据查询结果
                """,
                func=db_agent_executor.invoke
            )
        ]

        prompt = hub.pull("hwchase17/structured-chat-agent")

        agent = create_structured_chat_agent(
            llm=model,
            prompt=prompt,
            tools=tools
        )

        memory = ConversationBufferMemory(
            return_messages=True,
            memory_key="chat_history"
        )

        agent_executor = AgentExecutor.from_agent_and_tools(
            agent=agent,
            tools=tools,
            memory=memory,
            verbose=True,
            handle_parsing_errors=True
        )

        return agent_executor

    react_agent = create_react_agent(db_agent_executor)
    return react_agent.invoke({"input": input_text})
